package org.d8s.magicapi.plugin.interceptor;


import com.google.common.util.concurrent.RateLimiter;
import org.d8s.magicapi.util.HttpUtil;
import org.slf4j.Logger;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Component;
import org.ssssssss.magicapi.interceptor.RequestInterceptor;
import org.ssssssss.magicapi.model.ApiInfo;
import org.ssssssss.magicapi.model.JsonBean;
import org.ssssssss.script.MagicScriptContext;

import javax.annotation.PostConstruct;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import static org.d8s.magicapi.util.EncryptConstants.START_PLUGIN_LOG_MSG;
import static org.slf4j.LoggerFactory.getLogger;

/**
 * 全局接口请求频率限制
 * API接口访问上限,当访问频率或者并发量超过阈值应该禁止
 * Guava提供了限流工具类RateLimiter,该类基于令牌桶算法(Token Bucket)来完成限流
 * 创建一个稳定输出令牌的RateLimiter，保证了每个接口平均每秒不超过permitsPerSecond个请求
 * @author 冰点
 * @Date 2021-5-20 10:54:30
 */
@Component
@ConditionalOnProperty(prefix = "magic-plugin.limiter", name = "enable", havingValue = "true", matchIfMissing = false)
public class GlobalApiLimiterInterceptor implements RequestInterceptor {
    private static final Logger log = getLogger(GlobalApiLimiterInterceptor.class);
    private static RateLimiter rateLimiter;
    /**
     * 令牌速率
     */
    @Value("${magic-plugin.limiter.permits.per.second:20}")
    private double permitsPerSecond;
    /**
     * 预热期(warmup period)内，RateLimiter会平滑的将其释放令牌的速率加大，直到起达到最大速率
     */
    @Value("${magic-plugin.limiter.permits.warmup.Period:1}")
    private double warmupPeriod;
    /**
     * 预热期单位
     */
    @Value("${magic-plugin.limiter.permits.time.unit:MINUTES}")
    private TimeUnit timeUnit;
    /**
     * 获取令牌的最大等待时间
     */
    @Value("${magic-plugin.limiter.acquire.timeout:3}")
    private long timeout;
    private static Map<String, RateLimiter> apiRateLimiterMap;

    @PostConstruct
    public void initGlobalApiLimiterInterceptor() {
        log.info(START_PLUGIN_LOG_MSG,"全局接口请求频率","magic-plugin.limiter.enable=false","magic-plugin.limiter.permits.per.second:20");
        rateLimiter = RateLimiter.create(permitsPerSecond, 1, timeUnit);
        apiRateLimiterMap = new ConcurrentHashMap(8) {
            @Override
            public Object getOrDefault(Object key, Object defaultValue) {
                return super.getOrDefault(key, rateLimiter);
            }
        };
    }

    /**
     * 接口请求之前
     *
     * @param info    接口信息
     * @param context 脚本变量信息
     */
    @Override
    public Object preHandle(ApiInfo info, MagicScriptContext context, HttpServletRequest request, HttpServletResponse response) throws Exception {
        String path = info.getPath();
        String[] apiRateLimiterParams = info.getOptionValue("rateLimiter").split("/");
        String remoteHost = HttpUtil.getIpAddr(request);
        if (apiRateLimiterParams.length == 2 && !apiRateLimiterMap.containsKey(path) && !apiRateLimiterMap.containsKey(path)) {
            RateLimiter rateLimiter = RateLimiter.create(Double.valueOf(apiRateLimiterParams[0]), 1, TimeUnit.valueOf(apiRateLimiterParams[1]));
            apiRateLimiterMap.put(path, rateLimiter);
        }
        RateLimiter rateLimiter = apiRateLimiterMap.getOrDefault(path, null);
        if (!rateLimiter.tryAcquire(1, timeout, TimeUnit.SECONDS)) {
            log.error("[IP:{}]URL:[{}]请求超过频率限制,建议根据系统负载和数据库压力进行适当调节", remoteHost, path);
            return new JsonBean<>(100, String.format("[IP:%s]请求URL:[%s]超过频率限制", remoteHost, path));
        }
        return null;
    }

    /**
     * 接口执行之后
     *
     * @param info    接口信息
     * @param context 变量信息
     * @param value   即将要返回到页面的值
     */
    @Override
    public Object postHandle(ApiInfo info, MagicScriptContext context, Object value, HttpServletRequest request, HttpServletResponse response) throws Exception {
        log.info("{} 执行完毕，返回结果:{}", info.getName(), value);
        return null;
    }


}
